import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from model import AlexNet


def test_data():
    data = FashionMNIST(root='./data',
                        train=False,
                        download=True,
                        transform=transforms.Compose([transforms.Resize(size=227), transforms.ToTensor()]))
    test_data = Data.DataLoader(dataset=data,
                                batch_size=1,
                                shuffle=True)
    return test_data


def test(model, test_data):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    test_acc = 0.0
    test_num = 0

    classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    count = 0
    # 不计算梯度
    with torch.no_grad():
        for x, y in test_data:
            x = x.to(device)
            y = y.to(device)
            model.eval()
            output = model(x)
            pre_lab = torch.argmax(output, dim=1)
            result = pre_lab.item()
            label = y.item()
            test_acc += torch.sum(pre_lab == y.data)
            test_num += x.size(0)
            if result != label:
                count += 1
                print(f"错误预测{count} : 预测值：{classes[result]} ---------- 真实值：{classes[label]}")

    test_acc = test_acc / test_num * 100
    print("Test Accuracy: {:.2f}%".format(test_acc))


if __name__ == '__main__':
    model = AlexNet()
    # 加载模型参数
    model.load_state_dict(torch.load('checkpoints/AlexNet_epoch_20.pth'))
    test_data = test_data()
    test(model, test_data)

